{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FFI 数据类型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytest\n",
    "import pickle\n",
    "import numpy as np\n",
    "from tvm import ffi as tvm_ffi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "float32 = tvm_ffi.dtype(\"float32\")\n",
    "assert float32.__repr__() == \"dtype('float32')\"\n",
    "assert type(float32) == tvm_ffi.dtype\n",
    "x = np.array([1, 2, 3], dtype=float32)\n",
    "assert x.dtype == float32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "int32: 4\n",
      "float32x4: 16\n",
      "float8_e5m2x4: 4\n",
      "uint8: 1\n"
     ]
    }
   ],
   "source": [
    "for dtype_str in [\"int32\", \"float32x4\", \"float8_e5m2x4\", \"uint8\"]:\n",
    "    dtype = tvm_ffi.dtype(dtype_str)\n",
    "    print(f\"{dtype_str}: {dtype.itemsize}\")\n",
    "    # pickle\n",
    "    dtype_pickled = pickle.loads(pickle.dumps(dtype))\n",
    "    assert dtype_pickled.type_code == dtype.type_code\n",
    "    assert dtype_pickled.bits == dtype.bits\n",
    "    assert dtype_pickled.lanes == dtype.lanes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = tvm_ffi.dtype(\"float32\")\n",
    "dtype_with_lanes = dtype.with_lanes(4)\n",
    "assert dtype_with_lanes.type_code == dtype.type_code\n",
    "assert dtype_with_lanes.bits == dtype.bits\n",
    "assert dtype_with_lanes.lanes == 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
